﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;

namespace Luminal
{
    public enum MemoryAccessFlags : ulong
    {
        ReadOnly = cl_mem_flags.CL_MEM_READ_ONLY,
        WriteOnly = cl_mem_flags.CL_MEM_WRITE_ONLY,
        ReadWrite = cl_mem_flags.CL_MEM_READ_WRITE,
    }

    public enum MemoryObjectType : uint
    {
        Buffer = 0x10F0,
        Image2D = 0x10F1,
        Image3D = 0x10F2,
    }

    public enum ChannelOrder : uint
    {
        R = 0x10B0,
        A = 0x10B1,
        RG = 0x10B2,
        RA = 0x10B3,
        RGB = 0x10B4,
        RGBA = 0x10B5,
        BGRA = 0x10B6,
        ARGB = 0x10B7,
        Intensity = 0x10B8,
        Luminance = 0x10B9,
    }

    public enum ChannelType : uint
    {
        SignedNormalizedByte = 0x10D0,
        SignedNormalizedShort = 0x10D1,
        UnsignedNormalizedByte = 0x10D2,
        UnsignedNormalizedShort = 0x10D3,
        UnsignedNormalizedShort565 = 0x10D4,
        UnsignedNormalizedShort55 = 0x10D5,
        UnsignedNormalizedInteger101010 = 0x10D6,
        SignedByte = 0x10D7,
        SignedShort = 0x10D8,
        SignedInteger = 0x10D9,
        UnsignedByte = 0x10DA,
        UnsignedShort = 0x10DB,
        UnisgnedInteger = 0x10DC,
        HalfSingle = 0x10DD,
        Single = 0x10DE,
    }

    public struct ImageFormat
    {
        public ChannelOrder ChannelOrder;
        public ChannelType ChannelDataType;

        public ImageFormat(ChannelOrder channelOrder, ChannelType channelDataType)
        {
            ChannelOrder = channelOrder;
            ChannelDataType = channelDataType;
        }

        public static ImageFormat R32 { get { return new ImageFormat(ChannelOrder.R, ChannelType.Single); } }
        public static ImageFormat R8G8B8A8 { get { return new ImageFormat(ChannelOrder.RGBA, ChannelType.UnsignedNormalizedByte); } }
        public static ImageFormat R32G32B32A32 { get { return new ImageFormat(ChannelOrder.RGBA, ChannelType.Single); } }
        public static ImageFormat A8R8G8B8 { get { return new ImageFormat(ChannelOrder.ARGB, ChannelType.UnsignedNormalizedByte); } }
        public static ImageFormat A32R32G32B32 { get { return new ImageFormat(ChannelOrder.ARGB, ChannelType.Single); } }
    }

    public sealed class MemoryObject : LuminalObject
    {
        #region Interop
        [DllImport("OpenCL.dll")]
        static extern IntPtr clCreateBuffer(IntPtr context,
            cl_mem_flags flags,
            SizeT size,
            IntPtr host_ptr,
            out ErrorCode errcode_ret);

        [DllImport("OpenCL.dll")]
        static extern IntPtr clCreateImage2D(IntPtr context,
            cl_mem_flags flags,
            ref ImageFormat image_format,
            SizeT image_width,
            SizeT image_height,
            SizeT image_row_pitch,
            IntPtr host_ptr,
            out ErrorCode errcode_ret);

        [DllImport("OpenCL.dll")]
        static extern IntPtr clCreateImage3D(IntPtr context,
            cl_mem_flags flags,
            ref ImageFormat image_format,
            SizeT image_width,
            SizeT image_height,
            SizeT image_depth,
            SizeT image_row_pitch,
            SizeT image_slice_pitch,
            IntPtr host_ptr,
            out ErrorCode errcode_ret);

        [DllImport("OpenCL.dll")]
        static extern ErrorCode clReleaseMemObject(IntPtr memobj);

        [DllImport("OpenCL.dll")]
        static extern ErrorCode clGetMemObjectInfo(IntPtr memobj,
            cl_mem_info param_name,
            SizeT param_value_size,
            IntPtr param_value,
            out SizeT param_value_size_ret);

        [DllImport("OpenCL.dll")]
        static extern ErrorCode clGetImageInfo(IntPtr image,
            cl_image_info param_name,
            SizeT param_value_size,
            IntPtr param_value,
            out SizeT param_value_size_ret); 
        #endregion

        #region Fields
        int width, height, depth;
        int elementSize;
        ImageFormat format;
        #endregion

        #region Properties
        public Context Context { get; private set; }
        public int Size { get; private set; }
        public MemoryAccessFlags AccessFlags { get; private set; }
        public MemoryObjectType Type { get; private set; }
        public int Width
        {
            get
            {
                if (Type == MemoryObjectType.Buffer)
                    throw new InvalidOperationException("Memory object is not an image.");

                return width;
            }
        }
        public int Height
        {
            get
            {
                if (Type == MemoryObjectType.Buffer)
                    throw new InvalidOperationException("Memory object is not an image.");

                return height;
            }
        }
        public int Depth
        {
            get
            {
                if (Type != MemoryObjectType.Image3D)
                    throw new InvalidOperationException("Memory object is not a 3D image.");

                return depth;
            }
        }
        public int ElementSize
        {
            get
            {
                if (Type == MemoryObjectType.Buffer)
                    throw new InvalidOperationException("Memory object is not an image.");

                return elementSize;
            }
        }
        public ImageFormat Format
        {
            get
            {
                if (Type == MemoryObjectType.Buffer)
                    throw new InvalidOperationException("Memory object is not an image.");

                return format;
            }
        }

        private T GetInfo<T>(cl_mem_info info) where T : struct
        {
            SizeT infoSize = 0;
            Util.ThrowError("Could not get memory object property.", clGetMemObjectInfo(InternalPointer, info, 0, IntPtr.Zero, out infoSize));

            T output = default(T);
            Util.Pin(ref output, ptr =>
            {
                Util.ThrowError("Could not get memory object property.", clGetMemObjectInfo(InternalPointer, info, infoSize, ptr, out infoSize));
            });

            return output;
        }

        private T GetInfo<T>(cl_image_info info) where T : struct
        {
            SizeT infoSize = 0;
            Util.ThrowError("Could not get image property.", clGetImageInfo(InternalPointer, info, 0, IntPtr.Zero, out infoSize));

            T output = default(T);
            Util.Pin(ref output, ptr =>
            {
                Util.ThrowError("Could not get image property.", clGetImageInfo(InternalPointer, info, infoSize, ptr, out infoSize));
            });

            return output;
        }
        #endregion

        #region Constructor
        private MemoryObject(IntPtr ptr)
        {
            Initialize(ptr);

            try
            {
                Context = Context.FromPointer(GetInfo<IntPtr>(cl_mem_info.CL_MEM_CONTEXT));
                Size = (int)GetInfo<uint>(cl_mem_info.CL_MEM_SIZE);

                cl_mem_flags flags = (cl_mem_flags)GetInfo<ulong>(cl_mem_info.CL_MEM_FLAGS);
                if ((flags & cl_mem_flags.CL_MEM_READ_ONLY) == cl_mem_flags.CL_MEM_READ_ONLY)
                    AccessFlags = MemoryAccessFlags.ReadOnly;
                else if ((flags & cl_mem_flags.CL_MEM_WRITE_ONLY) == cl_mem_flags.CL_MEM_WRITE_ONLY)
                    AccessFlags = MemoryAccessFlags.WriteOnly;
                else
                    AccessFlags = MemoryAccessFlags.ReadWrite;

                Type = (MemoryObjectType)GetInfo<uint>(cl_mem_info.CL_MEM_TYPE);

                if (Type != MemoryObjectType.Buffer)
                {
                    elementSize = (int)GetInfo<uint>(cl_image_info.CL_IMAGE_ELEMENT_SIZE);
                    format = GetInfo<ImageFormat>(cl_image_info.CL_IMAGE_FORMAT);

                    width = (int)GetInfo<uint>(cl_image_info.CL_IMAGE_WIDTH);
                    height = (int)GetInfo<uint>(cl_image_info.CL_IMAGE_HEIGHT);

                    if (Type == MemoryObjectType.Image3D)
                        depth = (int)GetInfo<uint>(cl_image_info.CL_IMAGE_DEPTH);
                }
            }
            catch (Exception ex)
            {
                Dispose();
                throw ex;
            }
        }
        #endregion

        #region Buffer Creation
        public static MemoryObject CreateBuffer(Context context, MemoryAccessFlags accessFlags, int bufferSize)
        {
            if (context == null)
                throw new ArgumentNullException("context");
            if (context.Disposed)
                throw new ObjectDisposedException("context");
            if (bufferSize <= 0)
                throw new ArgumentOutOfRangeException("bufferSize", bufferSize, "The specified buffer size must be greater than zero.");
            if (context.Devices.Any(n => bufferSize > n.MaximumMemoryAllocationSize))
                throw new ArgumentOutOfRangeException("bufferSize", bufferSize, "The specified buffer size is too large for one or more of the devices.");

            cl_mem_flags flags = (cl_mem_flags)accessFlags | cl_mem_flags.CL_MEM_ALLOC_HOST_PTR;

            ErrorCode errorCode;
            IntPtr outputPtr = clCreateBuffer(context.InternalPointer, flags, bufferSize, IntPtr.Zero, out errorCode);

            Util.ThrowError(
                "Could not create buffer.", 
                errorCode,
                new ErrorCase(ErrorCode.CL_MEM_OBJECT_ALLOCATION_FAILURE, new MemoryObjectAllocationFailureException("Failed to allocate memory for the buffer object.")));

            return new MemoryObject(outputPtr);
        }

        public static MemoryObject CreateBuffer(Context context, MemoryAccessFlags accessFlags, Type elementType, int elementCount)
        {
            if (elementType == null)
                throw new ArgumentNullException("elementType");
            if (elementCount <= 0)
                throw new ArgumentOutOfRangeException("elementCount", elementCount, "The specified amount of elements must be greater than zero.");

            return CreateBuffer(context, accessFlags, elementCount * Marshal.SizeOf(elementType));
        }

        public static MemoryObject CreateBuffer<T>(Context context, MemoryAccessFlags accessFlags, int elementCount) where T : struct
        {
            return CreateBuffer(context, accessFlags, typeof(T), elementCount);
        }

        public static MemoryObject CreateBuffer<T>(Context context, MemoryAccessFlags accessFlags, T[] data) where T : struct
        {
            if (context == null)
                throw new ArgumentNullException("context");
            if (context.Disposed)
                throw new ObjectDisposedException("context");
            if (data == null)
                throw new ArgumentNullException("data");

            cl_mem_flags flags = (cl_mem_flags)accessFlags | cl_mem_flags.CL_MEM_COPY_HOST_PTR;

            int dataSize = data.Length * Marshal.SizeOf(typeof(T));

            ErrorCode errorCode = ErrorCode.CL_SUCCESS;
            IntPtr outputPointer = IntPtr.Zero;
            Util.Pin(data, ptr =>
            {
                outputPointer = clCreateBuffer(context.InternalPointer, flags, dataSize, ptr, out errorCode);
            });

            Util.ThrowError("Could not create buffer.", errorCode);

            MemoryObject output = new MemoryObject(outputPointer);

            return output;
        }
        #endregion

        #region Image2D Creation
        public static MemoryObject CreateImage2D(Context context, MemoryAccessFlags accessFlags, ImageFormat format, int width, int height)
        {
            if (context == null)
                throw new ArgumentNullException("context");
            if (context.Disposed)
                throw new ObjectDisposedException("context");
            if (width <= 0)
                throw new ArgumentOutOfRangeException("width", width, "The specified width must be greater than zero.");
            if (height <= 0)
                throw new ArgumentOutOfRangeException("height", height, "The specified height must be greater than zero.");
            if (context.Devices.Any(n => width > n.MaximumImage2DWidth))
                throw new ArgumentOutOfRangeException("width", width, "The specified width is too large for one or more of the devices.");
            if (context.Devices.Any(n => height > n.MaximumImage2DHeight))
                throw new ArgumentOutOfRangeException("height", height, "The specified height is too large for one or more of the devices.");

            cl_mem_flags flags = (cl_mem_flags)accessFlags | cl_mem_flags.CL_MEM_ALLOC_HOST_PTR;

            ErrorCode errorCode;
            IntPtr outputPtr = clCreateImage2D(context.InternalPointer, flags, ref format, width, height, 0, IntPtr.Zero, out errorCode);

            Util.ThrowError(
                "Could not create image.",
                errorCode,
                new ErrorCase(ErrorCode.CL_MEM_OBJECT_ALLOCATION_FAILURE, new MemoryObjectAllocationFailureException("Failed to allocate memory for the buffer object.")));

            return new MemoryObject(outputPtr);
        }

        public static MemoryObject CreateImage2D<T>(Context context, MemoryAccessFlags accessFlags, ImageFormat format, T[,] data)
        {
            if (data == null)
                throw new ArgumentNullException("data");
            
            int width = data.GetLength(0), height = data.GetLength(1);
            T[] flatData = new T[width * height];
            for (int i = 0; i < flatData.Length; i++)
                flatData[i] = data[i % width, i / width];

            return CreateImage2D<T>(context, accessFlags, format, width, height, flatData);
        }

        public static MemoryObject CreateImage2D<T>(Context context, MemoryAccessFlags accessFlags, ImageFormat format, int width, int height, T[] data)
        {
            if (context == null)
                throw new ArgumentNullException("context");
            if (context.Disposed)
                throw new ObjectDisposedException("context");
            if (width <= 0)
                throw new ArgumentOutOfRangeException("width", width, "The specified width must be greater than zero.");
            if (height <= 0)
                throw new ArgumentOutOfRangeException("height", height, "The specified height must be greater than zero.");
            if (context.Devices.Any(n => width > n.MaximumImage2DWidth))
                throw new ArgumentOutOfRangeException("width", width, "The specified width is too large for one or more of the devices.");
            if (context.Devices.Any(n => height > n.MaximumImage2DHeight))
                throw new ArgumentOutOfRangeException("height", height, "The specified height is too large for one or more of the devices.");
            if (data == null)
                throw new ArgumentNullException("data");

            cl_mem_flags flags = (cl_mem_flags)accessFlags | cl_mem_flags.CL_MEM_COPY_HOST_PTR;

            IntPtr outputPtr = IntPtr.Zero;
            ErrorCode errorCode = ErrorCode.CL_SUCCESS;
            Util.Pin(data, ptr =>
            {
                outputPtr = clCreateImage2D(context.InternalPointer, flags, ref format, width, height, 0, ptr, out errorCode);
            });

            Util.ThrowError(
                "Could not create image.",
                errorCode,
                new ErrorCase(ErrorCode.CL_MEM_OBJECT_ALLOCATION_FAILURE, new MemoryObjectAllocationFailureException("Failed to allocate memory for the buffer object.")));

            return new MemoryObject(outputPtr);
        }
        #endregion

        #region Image3D Creation
        public static MemoryObject CreateImage3D(Context context, MemoryAccessFlags accessFlags, ImageFormat format, int width, int height, int depth)
        {
            if (context == null)
                throw new ArgumentNullException("context");
            if (context.Disposed)
                throw new ObjectDisposedException("context");
            if (width <= 0)
                throw new ArgumentOutOfRangeException("width", width, "The specified width must be greater than zero.");
            if (height <= 0)
                throw new ArgumentOutOfRangeException("height", height, "The specified height must be greater than zero.");
            if (depth <= 0)
                throw new ArgumentOutOfRangeException("depth", depth, "The specified depth must be greater than zero.");
            if (context.Devices.Any(n => width > n.MaximumImage3DWidth))
                throw new ArgumentOutOfRangeException("width", width, "The specified width is too large for one or more of the devices.");
            if (context.Devices.Any(n => height > n.MaximumImage3DHeight))
                throw new ArgumentOutOfRangeException("height", height, "The specified height is too large for one or more of the devices.");
            if (context.Devices.Any(n => depth > n.MaximumImage3DDepth))
                throw new ArgumentOutOfRangeException("depth", depth, "The specified depth is too large for one or more of the devices.");

            cl_mem_flags flags = (cl_mem_flags)accessFlags | cl_mem_flags.CL_MEM_ALLOC_HOST_PTR;

            ErrorCode errorCode;
            IntPtr outputPtr = clCreateImage3D(context.InternalPointer, flags, ref format, width, height, depth, 0, 0, IntPtr.Zero, out errorCode);

            Util.ThrowError(
                "Could not create image.",
                errorCode,
                new ErrorCase(ErrorCode.CL_MEM_OBJECT_ALLOCATION_FAILURE, new MemoryObjectAllocationFailureException("Failed to allocate memory for the buffer object.")));

            return new MemoryObject(outputPtr);
        }

        public static MemoryObject CreateImage3D<T>(Context context, MemoryAccessFlags accessFlags, ImageFormat format, T[,,] data)
        {
            if (data == null)
                throw new ArgumentNullException("data");

            int width = data.GetLength(0), height = data.GetLength(1), depth = data.GetLength(2);
            int wh = width * height;

            T[] flatData = new T[width * height * depth];
            for (int i = 0; i < flatData.Length; i++)
            {
                flatData[i] = data[(i % wh) % width, (i % wh) / width, i / wh];
            }

            return CreateImage3D<T>(context, accessFlags, format, width, height, depth, flatData);
        }

        public static MemoryObject CreateImage3D<T>(Context context, MemoryAccessFlags accessFlags, ImageFormat format, int width, int height, int depth, T[] data)
        {
            if (context == null)
                throw new ArgumentNullException("context");
            if (context.Disposed)
                throw new ObjectDisposedException("context");
            if (width <= 0)
                throw new ArgumentOutOfRangeException("width", width, "The specified width must be greater than zero.");
            if (height <= 0)
                throw new ArgumentOutOfRangeException("height", height, "The specified height must be greater than zero.");
            if (depth <= 0)
                throw new ArgumentOutOfRangeException("depth", depth, "The specified depth must be greater than zero.");
            if (context.Devices.Any(n => width > n.MaximumImage3DWidth))
                throw new ArgumentOutOfRangeException("width", width, "The specified width is too large for one or more of the devices.");
            if (context.Devices.Any(n => height > n.MaximumImage3DHeight))
                throw new ArgumentOutOfRangeException("height", height, "The specified height is too large for one or more of the devices.");
            if (context.Devices.Any(n => depth > n.MaximumImage3DDepth))
                throw new ArgumentOutOfRangeException("depth", depth, "The specified depth is too large for one or more of the devices.");
            if (data == null)
                throw new ArgumentNullException("data");

            cl_mem_flags flags = (cl_mem_flags)accessFlags | cl_mem_flags.CL_MEM_COPY_HOST_PTR;

            IntPtr outputPtr = IntPtr.Zero;
            ErrorCode errorCode = ErrorCode.CL_SUCCESS;
            Util.Pin(data, ptr =>
            {
                outputPtr = clCreateImage3D(context.InternalPointer, flags, ref format, width, height, depth, 0, 0, ptr, out errorCode);
            });

            Util.ThrowError(
                "Could not create image.",
                errorCode,
                new ErrorCase(ErrorCode.CL_MEM_OBJECT_ALLOCATION_FAILURE, new MemoryObjectAllocationFailureException("Failed to allocate memory for the buffer object.")));

            return new MemoryObject(outputPtr);
        }
        #endregion

        protected override void Dispose(bool disposing)
        {
            if (disposing)
                clReleaseMemObject(InternalPointer);

            base.Dispose(disposing);
        }

        public static MemoryObject FromPointer(IntPtr ptr)
        {
            return ObjectTracker.ContainsObject(ptr) ? ObjectTracker.GetObject<MemoryObject>(ptr) : new MemoryObject(ptr);
        }
    }
}
